package com.boot.admin.modules.sys.ai.service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.Base64;

/**
 * 流式WebSocket服务 - 处理音频流传输
 */
@Slf4j
@Service
public class StreamingWebSocketService implements WebSocketHandler {
    
    // 存储WebSocket会话，按对话ID分组
    private final ConcurrentHashMap<String, CopyOnWriteArrayList<WebSocketSession>> conversationSessions 
        = new ConcurrentHashMap<>();
    
    // 存储会话与对话ID的映射
    private final ConcurrentHashMap<String, String> sessionConversationMap = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("WebSocket连接建立: {}", session.getId());
        
        // 从查询参数中获取对话ID
        String conversationId = extractConversationId(session);
        if (conversationId != null) {
            // 将会话添加到对应的对话组中
            conversationSessions.computeIfAbsent(conversationId, k -> new CopyOnWriteArrayList<>())
                                .add(session);
            sessionConversationMap.put(session.getId(), conversationId);
            
            log.info("WebSocket会话 {} 加入对话 {}", session.getId(), conversationId);
        } else {
            log.warn("WebSocket连接缺少对话ID参数");
            session.close();
        }
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        if (message instanceof TextMessage) {
            handleTextMessage(session, (TextMessage) message);
        } else if (message instanceof BinaryMessage) {
            handleBinaryMessage(session, (BinaryMessage) message);
        }
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误: {}", session.getId(), exception);
        removeSession(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        log.info("WebSocket连接关闭: {}, 状态: {}", session.getId(), closeStatus);
        removeSession(session);
    }

    @Override
    public boolean supportsPartialMessages() {
        return true;
    }

    /**
     * 处理文本消息
     */
    private void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
        String payload = message.getPayload();
        log.debug("接收到文本消息: {}", payload);
        
        // 这里可以处理客户端发送的控制消息
        // 例如：{"action": "start_audio"} 或 {"action": "stop_audio"}
        
        // 回显确认消息
        session.sendMessage(new TextMessage("{\"type\":\"ack\",\"message\":\"received\"}"));
    }

    /**
     * 处理二进制消息（语音数据）
     */
    private void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws IOException {
        byte[] audioData = message.getPayload().array();
        String conversationId = sessionConversationMap.get(session.getId());
        
        log.debug("接收到语音数据: {} bytes, 对话ID: {}", audioData.length, conversationId);
        
        // 这里可以处理客户端发送的语音数据，例如：
        // 1. 实时语音识别
        // 2. 语音流式处理
        // 目前只是记录日志
    }

    /**
     * 发送音频chunk给指定对话的所有客户端
     */
    public void sendAudioChunk(String conversationId, String messageId, byte[] audioChunk) {
        CopyOnWriteArrayList<WebSocketSession> sessions = conversationSessions.get(conversationId);
        if (sessions == null || sessions.isEmpty()) {
            log.debug("对话 {} 没有活跃的WebSocket连接", conversationId);
            return;
        }

        // 构建音频消息
        String audioMessage = String.format(
            "{\"type\":\"audio_chunk\",\"conversationId\":\"%s\",\"messageId\":\"%s\",\"data\":\"%s\"}", 
            conversationId, messageId, Base64.getEncoder().encodeToString(audioChunk)
        );

        // 发送给该对话的所有连接
        sessions.removeIf(session -> {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(audioMessage));
                    log.debug("发送音频chunk给会话 {}: {} bytes", session.getId(), audioChunk.length);
                    return false; // 保留该会话
                } else {
                    log.debug("会话 {} 已关闭，将被移除", session.getId());
                    return true; // 移除已关闭的会话
                }
            } catch (Exception e) {
                log.error("发送音频chunk失败，会话: {}", session.getId(), e);
                return true; // 移除出错的会话
            }
        });
    }

    /**
     * 发送音频流结束信号
     */
    public void sendAudioComplete(String conversationId, String messageId) {
        CopyOnWriteArrayList<WebSocketSession> sessions = conversationSessions.get(conversationId);
        if (sessions == null || sessions.isEmpty()) {
            return;
        }

        String completeMessage = String.format(
            "{\"type\":\"audio_complete\",\"conversationId\":\"%s\",\"messageId\":\"%s\"}", 
            conversationId, messageId
        );

        sessions.removeIf(session -> {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(completeMessage));
                    log.debug("发送音频完成信号给会话 {}", session.getId());
                    return false;
                } else {
                    return true;
                }
            } catch (Exception e) {
                log.error("发送音频完成信号失败，会话: {}", session.getId(), e);
                return true;
            }
        });
    }

    /**
     * 发送错误消息
     */
    public void sendError(String conversationId, String error) {
        CopyOnWriteArrayList<WebSocketSession> sessions = conversationSessions.get(conversationId);
        if (sessions == null || sessions.isEmpty()) {
            return;
        }

        String errorMessage = String.format(
            "{\"type\":\"error\",\"conversationId\":\"%s\",\"error\":\"%s\"}", 
            conversationId, error
        );

        sessions.removeIf(session -> {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(errorMessage));
                    log.debug("发送错误消息给会话 {}: {}", session.getId(), error);
                    return false;
                } else {
                    return true;
                }
            } catch (Exception e) {
                log.error("发送错误消息失败，会话: {}", session.getId(), e);
                return true;
            }
        });
    }

    /**
     * 从WebSocket会话中提取对话ID
     */
    private String extractConversationId(WebSocketSession session) {
        try {
            String query = session.getUri().getQuery();
            if (query != null && query.contains("conversationId=")) {
                String[] params = query.split("&");
                for (String param : params) {
                    if (param.startsWith("conversationId=")) {
                        return param.split("=")[1];
                    }
                }
            }
        } catch (Exception e) {
            log.error("提取对话ID失败", e);
        }
        return null;
    }

    /**
     * 移除会话
     */
    private void removeSession(WebSocketSession session) {
        String conversationId = sessionConversationMap.remove(session.getId());
        if (conversationId != null) {
            CopyOnWriteArrayList<WebSocketSession> sessions = conversationSessions.get(conversationId);
            if (sessions != null) {
                sessions.remove(session);
                if (sessions.isEmpty()) {
                    conversationSessions.remove(conversationId);
                }
            }
            log.info("WebSocket会话 {} 从对话 {} 中移除", session.getId(), conversationId);
        }
    }

    /**
     * 获取对话的活跃连接数
     */
    public int getActiveConnectionCount(String conversationId) {
        CopyOnWriteArrayList<WebSocketSession> sessions = conversationSessions.get(conversationId);
        return sessions != null ? sessions.size() : 0;
    }

    /**
     * 获取所有活跃对话数
     */
    public int getActiveConversationCount() {
        return conversationSessions.size();
    }

    /**
     * 获取总连接数
     */
    public int getTotalConnectionCount() {
        return sessionConversationMap.size();
    }
}